"""
随机游走算法
用于在依赖图上进行根因定位的随机游走评分器
"""
from collections import defaultdict
from typing import Callable
from typing import Dict
from typing import List
from typing import Union

import numpy as np
import pandas as pd

from .base import Score
from .base import Scorer
from ..model.case import CaseData
from ..model.graph import Graph
from ..model.graph import Node


def _times(num: int, multiplier: int = 10) -> int:
    """计算随机游走的迭代次数，默认为节点数的10倍"""
    return num * multiplier


class RandomWalkScorer(Scorer):
    """
    基于随机游走的评分器
    通过在依赖图上进行随机游走来评估节点的重要性
    """

    def __init__(
        self,
        rho: float = 0.5,          # 子节点权重衰减因子
        remove_sli: bool = False,   # 是否移除SLI节点
        num_loop: Union[int, Callable[[int], int]] = None,  # 随机游走迭代次数
        **kwargs,
    ):
        """
        初始化随机游走评分器
        
        Args:
            rho: 子节点权重衰减因子，用于控制向下传播的权重
            remove_sli: 是否在计算转移矩阵时移除SLI节点
            num_loop: 随机游走迭代次数，可以是固定值或基于节点数的计算函数
        """
        super().__init__(**kwargs)
        self._rho = rho
        self._remove_sli = remove_sli
        self._num_loop = num_loop if num_loop is not None else _times
        self._rng = np.random.default_rng(self._seed)  # 随机数生成器

    def generate_transition_matrix(
        self, graph: Graph, data: CaseData, scores: Dict[Node, Score]
    ) -> pd.DataFrame:
        """
        生成随机游走的转移矩阵
        
        计算规则：
        1. 子节点权重 = rho * 子节点异常分数
        2. 父节点权重 = 父节点异常分数
        3. 自环权重 = max(自身异常分数 - 最大出边权重, 0)
        4. 所有权重归一化处理
        """
        # 初始化转移矩阵
        nodes = list(scores.keys())
        size = len(nodes)
        matrix = pd.DataFrame(np.zeros([size, size]), index=nodes, columns=nodes)
        
        # 为每个节点计算转移概率
        for node in scores:
            # 计算到子节点的转移概率
            for child in graph.children(node):
                if child in scores:
                    matrix[node][child] = self._rho * abs(scores[child].score)

            # 计算到父节点的转移概率
            parents = graph.parents(node)
            if self._remove_sli:  # 可选择性移除SLI节点
                parents -= {data.sli}
            for parent in parents:
                if parent in scores:
                    matrix[node][parent] = abs(scores[parent].score)

            # 计算自环概率
            matrix.loc[node, node] = max(abs(scores[node].score) - matrix.loc[node].max(), 0)
            # 归一化处理
            total_weight = matrix[node].sum()
            if total_weight > 0:
                matrix[node] /= total_weight
            else:  # 如果所有权重都为0，则均匀分布
                matrix[node] = 1 / size
        return matrix

    def _walk(
        self, start: Node, num_loop: int, matrix: pd.DataFrame
    ) -> Dict[Node, int]:
        """
        执行随机游走过程
        
        Args:
            start: 起始节点（通常是SLI节点）
            num_loop: 游走步数
            matrix: 转移矩阵
            
        Returns:
            Dict[Node, int]: 记录每个节点被访问的次数
        """
        node: Node = start
        counter = defaultdict(int)  # 访问计数器
        for _ in range(num_loop):
            # 根据转移概率选择下一个节点
            node = self._rng.choice(matrix.index, p=matrix[node])
            counter[node] += 1
        return counter

    def score(
        self,
        graph: Graph,
        data: CaseData,
        current: float,
        scores: Dict[Node, Score] = None,
    ) -> Dict[Node, Score]:
        """
        执行评分过程
        
        工作流程：
        1. 生成转移矩阵
        2. 执行随机游走
        3. 根据访问次数计算最终得分
        """
        if not scores:
            return scores

        # 生成转移矩阵
        matrix = self.generate_transition_matrix(graph=graph, data=data, scores=scores)
        # 确定迭代次数
        if isinstance(self._num_loop, int):
            num_loop = self._num_loop
        else:
            num_loop = self._num_loop(len(scores))
        # 执行随机游走
        counter = self._walk(start=data.sli, num_loop=num_loop, matrix=matrix)

        # 更新节点得分
        for node, score in scores.items():
            score["pagerank"] = score.score = counter[node] / num_loop
        return scores


class SecondOrderRandomWalkScorer(RandomWalkScorer):
    """
    基于二阶随机游走的评分器
    考虑前一步的转移概率来影响下一步的选择
    """

    def __init__(self, beta: float = 0.5, **kwargs):
        """
        初始化二阶随机游走评分器
        
        Args:
            beta: 平衡因子，用于权衡当前转移概率和前一步转移概率的影响
        """
        super().__init__(**kwargs)
        self._beta = beta

    def _walk(
        self, start: Node, num_loop: int, matrix: pd.DataFrame
    ) -> Dict[Node, int]:
        """
        执行二阶随机游走
        
        特点：
        1. 考虑前一步的转移概率
        2. 使用beta参数平衡两种概率的影响
        3. 更好地捕捉路径依赖特征
        """
        node: Node = start        # 当前节点
        node_pre: Node = start    # 前一个节点
        counter = defaultdict(int)  # 访问计数器
        
        for _ in range(num_loop):
            # 获取前一步的转移概率
            prob_pre = matrix[node_pre][node]
            node_pre = node
    
            # 计算候选节点和权重
            candidates: List[Node] = []
            weights: List[float] = []
            # 将 iteritems() 改为 items()
            for key, value in matrix[node].items():
                if value > 0:
                    candidates.append(key)
                    # 结合前一步和当前步的转移概率
                    weights.append((1 - self._beta) * prob_pre + self._beta * value)
                    
            # 选择下一个节点
            total_weight = sum(weights)
            if total_weight == 0:  # 如果所有权重为0，随机选择
                node = self._rng.choice(candidates)
            else:  # 根据权重概率选择
                node = self._rng.choice(
                    candidates, p=[weight / total_weight for weight in weights]
                )
            counter[node] += 1
        return counter
